'''
当前代码所读pcd格式为 x y z rgb label object
'''
import os
from itertools import islice


def get_label(pcd_name):
    '''
    获取pcd文件的label
    '''
    # 完整的标注文件
    pcd = open("./labeled/" + pcd_name, 'r')
    # 点云xyz坐标文件
    xyz = open("./xyz/" + pcd_name.replace(".pcd", ".pts"), 'a')
    # 标签文件
    label = open("./label/" + pcd_name.replace(".pcd", ".seg"), 'a')

    # pcd 的 header 一般为10行,从第11行开始读
    # islice(iterable, start, stop, step)
    # iterable(可迭代对象), start(开始迭代位置), stop(结束迭代位置), step(迭代步长)
    for line in islice(pcd, 10, None):
        # 检查是否到文件尾
        if not line:
            break
        item = line.split(" ")
        # 前三项xyz保存到点云文件
        xyz.write(item[0] + " " + item[1] + " " + item[2] + "\n")
        # 第5项,标签,保存到标签文件
        label.write(item[4] + "\n")
        # print(line.split(' '))
    pcd.close()
    label.close()

def seg_point(pcd_name):
    '''
    根据label分割object
    '''
    # 完整的标注文件
    pcd = open("./labeled/" + pcd_name, 'r')
    # 点云xyz坐标文件
    segment = open("./segment/" + pcd_name, 'a')

    # pcd 的 header 一般为10行,从第11行开始读
    # islice(iterable, start, stop, step)
    # iterable(可迭代对象), start(开始迭代位置), stop(结束迭代位置), step(迭代步长)
    for line in islice(pcd, 10, None):
        # 检查是否到文件尾
        if not line:
            break
        item = line.split(" ")
        if item[4] == '1':
            # 前三项xyz保存到点云文件
            segment.write(item[0] + " " + item[1] + " " + item[2] + "\n")

        # print(line.split(' '))
    pcd.close()
    segment.close()


if __name__ == "__main__":
    # 获取指定目录下的文件名存入列表
    pcds = os.listdir('./labeled')

    i = 0
    for pcd_name in pcds:
        print(pcd_name + "... ", end=" ")
        seg_point(pcd_name)
        print(" done.")
        i += 1
    print(str(i) + " files have done.")
